package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;


public class CryptoUtil {

    public static void main(String[] args) throws NoSuchAlgorithmException, IOException, InvalidKeySpecException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        /*
        Path pubPath = Paths.get("/home/caritas/dslab/keys/client/transfer-1_pub.der");
        Path privPath = Paths.get("/home/caritas/dslab/keys/server/transfer-1.der");
        PublicKey publicKey = getPublicKey(pubPath);
        PrivateKey privateKey = getPrivateKey(privPath);
        String message = "Ich mag Schnitzel :3";
        System.out.println(message);

        Cipher encCipher = createRSAEncryptCipher(publicKey);
        System.out.println(encCipher.doFinal(message.getBytes(StandardCharsets.UTF_8)).length);
        String encMessage64 = new String(Base64.getEncoder().encode(encCipher.doFinal(message.getBytes(StandardCharsets.UTF_8))));
        System.out.println(encMessage64.length());
        System.out.println(encMessage64);

        Cipher decCipher = createRSADecryptCipher(privateKey);
        String decMessage = new String(decCipher.doFinal(Base64.getDecoder().decode(encMessage64.getBytes(StandardCharsets.UTF_8))));
        System.out.println(decMessage);

        encCipher = createRSAEncryptCipher(publicKey);
        SecretKey symKey = generateNewKeydata().key();
        String keyString = encryptMessage(new String(symKey.getEncoded()), encCipher);

        decCipher = createRSADecryptCipher(privateKey);
        String encKeyString = decryptMessage(keyString, decCipher);
        System.out.println("abs " + new String(symKey.getEncoded()));
        System.out.println("abs " + encKeyString);
        */

/*
        String message = "Fuck Java";
        Cipher cipher = Cipher.getInstance("AES/CTR/NoPadding");
        Pair<SecretKey, byte[]> keyData = generateNewKeydata();
        cipher.init(Cipher.ENCRYPT_MODE, keyData.key(), new IvParameterSpec(keyData.value()));
        String enc = encryptMessage(message, cipher);
        String key = encodeB64(keyData.key().getEncoded());
        String iv = encodeB64(keyData.value());


        SecretKey secret = decodeSecret(key);
        cipher.init(Cipher.DECRYPT_MODE, secret, new IvParameterSpec(decodeB64(iv)));
        String messdec = decryptMessage(enc, cipher);
        System.out.println(messdec);
    */
        /*
        String message = "Fuck Java";
        Cipher cipher = Cipher.getInstance("AES/CTR/NoPadding");
        Pair<SecretKey, byte[]> keyData = generateNewKeydata();
        cipher.init(Cipher.ENCRYPT_MODE, keyData.key(), new IvParameterSpec(keyData.value()));
        String messageEnc = new String(cipher.doFinal(message.getBytes(StandardCharsets.UTF_8)));
        System.out.println(messageEnc);
        System.out.println(keyData.key().getEncoded().length);
        System.out.println(encodeB64(keyData.key().getEncoded()));
        String keyo = encodeB64(keyData.key().getEncoded());

        System.out.println(new String(decodeB64(encodeB64("abc".getBytes(StandardCharsets.UTF_8)))));
        System.out.println("pen " + Arrays.toString(keyData.value()));
        System.out.println(encodeB64(keyData.value()));
        System.out.println("pen " + Arrays.toString(decodeB64(encodeB64(keyData.value()))));
        IvParameterSpec ivspec = new IvParameterSpec(decodeB64(encodeB64(keyData.value())));
        System.out.println("pen " + Arrays.toString(ivspec.getIV()));
        cipher.init(Cipher.DECRYPT_MODE, decodeSecret(decodeB64(keyo)), new IvParameterSpec(decodeB64(encodeB64(keyData.value()))));
*/



        /*
        // make sure to use the right ALGORITHM for what you want to do (see text)
        Mac hMac = Mac.getInstance("HmacSHA256");
        hMac.init(secretKey);
        // MESSAGE is the message to sign in bytes
        Message mess = new Message();
        mess.data = ":3";
        mess.sender = "pipi@pupu.com";
        mess.recipients = "burg@junge.com";
        mess.subject = "^^";
        hMac.update(mess.allFields().getBytes(StandardCharsets.UTF_8));
        byte[] hash = hMac.doFinal();
        System.out.println(new String(hash));

        // computedHash is the HMAC of the received plaintext
        byte[] computedHash = hMac.doFinal();
        // receivedHash is the HMAC that was sent by the communication partner
        byte[] receivedHash = hash;
        boolean validHash = MessageDigest.isEqual(computedHash, receivedHash);
        System.out.println(validHash);
        */

        SecretKey secretKey = getSecretKey(Paths.get("/home/caritas/dslab/keys/hmac.key"));

        Message mess = new Message();
        mess.data = ":3";
        mess.sender = "pipi@pupu.com";
        mess.recipients = "burg@junge.com";
        mess.subject = "^^";

        mess.hash = hashMessage(mess, secretKey);
        System.out.println(verifyHash(mess, secretKey));
    }

    /**
     * Generates a new AES secret key.
     * @return AES Secret Key
     * @throws NoSuchAlgorithmException Should not happen, but may if library has problems.
     */
    public static SecretKey generateSecret() throws NoSuchAlgorithmException {
        KeyGenerator keyGenerator;
        try {
            keyGenerator = KeyGenerator.getInstance("AES");
        } catch (NoSuchAlgorithmException exception) {
            // Should not happen
            throw exception;
        }
        keyGenerator.init(256);
        return keyGenerator.generateKey();
    }

    /**
     * Creates an array of secure random numbers of specified size.
     * @param size Size of returned Array.
     * @return A byte array of secure random numbers.
     */
    public static byte[] generateRandom(int size) {
        byte[] rand = new byte[size];
        new SecureRandom().nextBytes(rand);
        return rand;
    }

    /**
     * Creates a Pair object containing a secret key (AES) and an initialisation vector (as a byte array).
     * @return Pair containing a secret key and a byte array of size 16
     * @throws NoSuchAlgorithmException - Should never happen, could in case of a library failure.
     */
    public static Pair<SecretKey, byte[]> generateNewKeydata() throws NoSuchAlgorithmException {
        return new Pair<>(generateSecret(), generateRandom(16));
    }

    /**
     * Recovers an AES secret key from a String.
     * @param key String containing an encoded secret key.
     * @return A secret key recovered from the specified String.
     */
    public static SecretKey decodeSecret(String key) {
        return new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), 0, key.length(), "AES");
    }

    /**
     * Recovers an AES secret key from a byte Array.
     * @param key String containing an encoded secret key.
     * @return A secret key recovered from the specified byte Array.
     */
    public static SecretKey decodeSecret(byte[] key) {
        return new SecretKeySpec(key, 0, key.length, "AES");
    }

    /**
     * Encodes the specified String in Base64 and returns it as a String.
     * @param message String to be encoded.
     * @return A String encoded in Base64
     */
    public static String encodeB64(String message) {
        return Base64.getEncoder().encodeToString(message.getBytes(StandardCharsets.UTF_8));
    }

    /**
     * Encodes the specified byte Array in Base64 and returns it as a String.
     * @param message Byte Array to be encoded.
     * @return A String encoded in Base64
     */
    public static String encodeB64(byte[] message) {
        return Base64.getEncoder().encodeToString(message);
    }

    /**
     * Decodes the specified String from Base64 and returns it as a byte Array.
     * @param message String to be decoded.
     * @return A decoded byte Array.
     */
    public static byte[] decodeB64(String message) {
        return Base64.getDecoder().decode(message);
    }

    /**
     * Decodes the specified byte Array from Base64 and returns it as a byte Array.
     * @param message String to be decoded.
     * @return A decoded byte Array.
     */
    public static byte[] decodeB64(byte[] message) {
        return Base64.getDecoder().decode(message);
    }

    /**
     * Reads the private key from a specified file encoded in the .der format.
     * @param location Location of a file containing a private key.
     * @return A private Key read from a file.
     * @throws IOException If an IO error occurs
     */
    public static PrivateKey getPrivateKey(Path location) throws IOException{
        return Keys.readPrivateKey(location.toFile());
    }

    /**
     * Reads the public key from a specified file encoded in the .der format.
     * @param location Location of a file containing a public key.
     * @return A public Key read from a file.
     * @throws IOException If an IO error occurs
     */
    public static PublicKey getPublicKey(Path location) throws IOException {
        return Keys.readPublicKey(location.toFile());
    }

    /**
     * Reads the secret key from a specified file.
     * @param location Location of a file containing a secret key.
     * @return A secret key read from a file.
     * @throws IOException If an IO error occurs
     */
    public static SecretKey getSecretKey(Path location) throws IOException {
        return Keys.readSecretKey(location.toFile());
    }

    /**
     * Creates a Cipher object for encrypting in RSA (asymmetric).
     * @param key Key used for encrypting.
     * @return A cipher initialised for encrypting.
     */
    public static Cipher createRSAEncryptCipher(PublicKey key) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.ENCRYPT_MODE, key);
            return cipher;
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }

    /**
     * Creates a Cipher object for decrypting in RSA (asymmetric).
     * @param key Key used for decrypting.
     * @return A cipher initialised for decrypting.
     */
    public static Cipher createRSADecryptCipher(PrivateKey key) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.DECRYPT_MODE, key);
            return cipher;
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }

    /**
     * Encrypts a specified String with a specified cipher and encodes the encrypted String in Base64.
     * @param message Message to be encrypted
     * @param cipher Cipher to be used for encryption
     * @return An encrypted String in Base64.
     * @throws IllegalBlockSizeException If the provided String does not match the block size of the cipher in case a block-cipher is used.
     * @throws BadPaddingException If the input String has incorrect padding for the used cipher.
     */
    public static String encryptMessage(String message, Cipher cipher) throws IllegalBlockSizeException, BadPaddingException {
        return new String(Base64.getEncoder().encode(cipher.doFinal(message.getBytes(StandardCharsets.UTF_8))));
    }

    /**
     * Decrypts a specified String in Base64 with a specified cipher.
     * @param message String in Base64 to be decrypted
     * @param cipher Cipher to be used for decryption
     * @return A decrypted String.
     * @throws IllegalBlockSizeException If the provided String does not match the block size of the cipher in case a block-cipher is used.
     * @throws BadPaddingException If the input String has incorrect padding for the used cipher.
     */
    public static String decryptMessage(String message, Cipher cipher) throws IllegalBlockSizeException, BadPaddingException {
        return new String(cipher.doFinal(Base64.getDecoder().decode(message.getBytes(StandardCharsets.UTF_8))));
    }

    /**
     * Hashes a Message object using the HmacSHA256 algorithm with the provided secret key.
     * @param message Message to be hashed.
     * @param key Key to be used for hashing.
     * @return A string containing the hash of the specified String encoded in Base64
     * @throws NoSuchAlgorithmException Should never happen, could in case of a library failure.
     * @throws InvalidKeyException If an invalid key is provided.
     */
    public static String hashMessage(Message message, SecretKey key) throws NoSuchAlgorithmException, InvalidKeyException {
        return encodeB64(hashMessageArr(message, key));
    }

    /**
     * Verifies a HmacSHA256 hash with the provided secret key.
     * @param message Message object containing both info and their respective hash.
     * @param key Key to be used for hashing.
     * @return True if the Message.hash object is equal to the hash of the message object.
     * @throws NoSuchAlgorithmException Should never happen, could in case of a library failure.
     * @throws InvalidKeyException If an invalid key is provided.
     */
    public static boolean verifyHash(Message message, SecretKey key) throws NoSuchAlgorithmException, InvalidKeyException {
        // computedHash is the HMAC of the received plaintext
        byte[] computedHash = hashMessageArr(message, key);
        // receivedHash is the HMAC that was sent by the communication partner
        byte[] receivedHash = decodeB64(message.hash);
        return MessageDigest.isEqual(computedHash, receivedHash);
    }

    private static byte[] hashMessageArr(Message message, SecretKey key) throws NoSuchAlgorithmException, InvalidKeyException {
        // make sure to use the right ALGORITHM for what you want to do (see text)
        Mac hMac = Mac.getInstance("HmacSHA256");
        hMac.init(key);
        // MESSAGE is the message to sign in bytes
        hMac.update(message.allFields().getBytes(StandardCharsets.UTF_8));
        return hMac.doFinal();
    }

}
